/*
 * Comm.java
 *
 * Version:
 *      $Id: Comm.java 27 2009-04-17 02:30:39Z PErickson87 $
 *
 * Copyright (c) 2009 Peter O. Erickson
 *
 * Permission is hereby granted, free of charge, to any person
 * obtaining a copy of this software and associated documentation
 * files (the "Software"), to deal in the Software without major
 * restriction, including without limitation the rights to use,
 * copy, modify, merge, publish, and/or distribute copies of, but
 * under no circumstances sublicense and/or sell, the Software,
 * and to permit persons to whom the Software is furnished to do
 * so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 * OTHER DEALINGS IN THE SOFTWARE.
 */

package edu.rit.poe.cluster;

import edu.rit.poe.util.ClassLoaderOIS;
import java.io.BufferedReader;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.EOFException;
import java.io.FileReader;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.BrokenBarrierException;
import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * Parallel cluster communicator.
 *
 * @author  Peter O. Erickson
 *
 * @version $Id: Comm.java 27 2009-04-17 02:30:39Z PErickson87 $
 */
public final class Comm {
    
    /** The "well-known" cluster port for intercommunication channels. */
    public static final int PORT = 4242;
    
    /** The number of nodes in the parallel cluster. */
    protected static int size;
    
    /** The rank of this particular node. */
    protected static int rank;
    
    /** The hostnames of all nodes in the cluster (indexed by rank). */
    protected static String[] nodehosts;
    
    /** The barrier for all nodes to await. */
    protected static volatile CyclicBarrier barrier;
    
    /** The barrier to wait at the end of the node execution. */
    private static CyclicBarrier endBarrier;
    
    /** An array of the proxies representing nodes in the cluster. */
    protected static Node[] nodes;
    
    /** The current network class loader. */
    protected static ClassLoader ncl;
    
    // initialize size and rank (for server node)
    static {
        size = 1;
        rank = 0;
    }
    
    /**
     * Private constructor to block instantiation.
     */
    private Comm() {
    }
    
    /**
     * Initializes the communicator.  After invoking this method, all code
     * henceforth will be run in parallel, on the requested number of nodes.
     * 
     * @param   args            the command line arguments of the application
     * 
     * @throws  IOException     if there is any fatal IO error
     */
    public static void init( String[] args ) throws IOException {
        if ( rank == 0 ) {
            // determine the class that ran me
            StackTraceElement[] ste = new Throwable().getStackTrace();
            String classname = ste[ ste.length - 1 ].getClassName();
            
            // run the classloader server
            final NetworkClassLoader.Server loaderServer =
                    new NetworkClassLoader.Server();
            new Thread( loaderServer ).start();
            //TODO: FIX
            ncl = Comm.class.getClassLoader();
            
            // load all available hosts for nodes
            List<String> nodehostslist = new ArrayList<String>();
            BufferedReader filein =
                    new BufferedReader( new FileReader( "parallelhosts" ) );
            String line = null;
            while ( ( line = filein.readLine() ) != null ) {
                nodehostslist.add( line );
            }
            nodehosts = new String[ 0 ];
            nodehosts = nodehostslist.toArray( nodehosts );
            
            // how many processes does this job want?
            size = nodehostslist.size() + 1;
            try {
                size = Integer.parseInt( System.getProperty( "P" ) );
            } catch ( Exception e ) {
                // ignore, use all available
            }
            if ( size > ( nodehosts.length + 1 ) ) {
                throw new IOException( "Not enough nodes for job." );
            }
            // initialize the local barrier
            barrier = new CyclicBarrier( size );
            endBarrier = new CyclicBarrier( size );
            // start once for root
            new Thread() {
                @Override
                public void run() {
                    try {
                        endBarrier.await();
                    } catch ( Exception e ) {
                        System.err.println( "Error waiting at end barrier." );
                        e.printStackTrace();
                    }
                    // kill the network classloader
                    loaderServer.kill();
                }
            }.start();
            
            nodes = new Node[ size ];
            // FOR EACH NODE
            for ( short nodeRank = 1; nodeRank < ( nodehosts.length + 1 );
                    nodeRank++ ) {
                String hostname = nodehosts[ nodeRank - 1 ].trim();
                Socket socket = null;
                try {
                    socket = new Socket( hostname, PORT );
                } catch ( IOException e ) {
                    // kill the classloader server, prepare to die cleanly
                    loaderServer.kill();
                    throw e;
                }
                nodes[ nodeRank ] = new Node( nodeRank, socket );
                
                DataOutputStream out =
                        new DataOutputStream( socket.getOutputStream() );
                ObjectOutputStream oos =
                        new ObjectOutputStream( socket.getOutputStream() );
                out.writeUTF( classname );
                out.writeInt( size );
                out.writeInt( nodeRank );
                oos.writeObject( args );
                oos.writeObject( nodehosts );
                
                // nodes will now setup full interconnectivity
            }
        } else {
            // setup interconnectivity
            short connected = 1;
            ServerSocket ss = new ServerSocket( PORT );
            
            // receive until we are connected to #rank nodes
            Socket socket = null;
            while ( connected < rank ) {
                socket = ss.accept();
                nodes[ connected ] = new Node( connected, socket );
                
                // increment to next node
                connected++;
            }
            ss.close();
            
            // own node is null
            nodes[ rank ] = null;
            connected++;
            
            // send your own connections
            for ( ; connected < size; connected++ ) {
                socket = new Socket( nodehosts[ connected - 1 ], PORT );
                nodes[ connected ] = new Node( connected, socket );
            }
        }
        
        // wait for all nodes to finish
        try {
            barrier();
        } catch ( Exception e ) {
            // ignore
        }
    }
    
    /**
     * Returns the number of nodes in the cluster.
     * 
     * @return  the size of the cluster
     */
    public static int size() {
        return size;
    }
    
    /**
     * Returns the rank of this particular node.
     * 
     * @return  this node's rank
     */
    public static int rank() {
        return rank;
    }
    
    /**
     * Executes a cyclic barrier, and waits until all nodes have invoked this
     * method.
     * 
     * @throws  InterruptedException    if this thread is interrupted while
     *                                  waiting at the barrier
     * @throws  BrokenBarrierException  if the barrier is broken while waiting
     */
    public synchronized static void barrier() throws InterruptedException,
            BrokenBarrierException {
        // send barrier code
        for ( Node node : nodes ) {
            if ( node != null ) {
                try {
                    node.out.writeByte( Node.Command.BARRIER );
                } catch ( Exception e ) {
                    System.err.println( "Error sending barrier code." );
                    e.printStackTrace( System.err );
                }
            }
        }
        
        // wait for all to call barrier
        barrier.await();
    }
    
    /**
     * Sends the specified object to the specified node, with the tag of
     * <tt>0</tt>.
     * 
     * @param   rank            the rank of the node to send the object to
     * @param   o               the object to be sent to the <tt>rank</tt> node
     * 
     * @throws  IOException     if there is any communication error while
     *                          sending
     */
    public static void send( int rank, Object o ) throws IOException {
        send( rank, o, 0 );
    }
    
    /**
     * Sends the specified object to the specified node, and the specified tag.
     * 
     * @param   rank            the rank of the node to send the object to
     * @param   o               the object to be sent to the <tt>rank</tt> node
     * @param   tag             the tag under which to associate the object
     * 
     * @throws  IOException     if there is any communication error while
     *                          sending
     */
    public static void send( int rank, Object o, int tag ) throws IOException {
        send( rank, o, tag, false );
    }
    
    /**
     * Helper method for sending an object for a either a receive or a gather.
     * 
     * @param   rank            the rank of the node to send the object to
     * @param   o               the object to be sent to the <tt>rank</tt> node
     * @param   tag             the tag under which to associate the object
     * @param   gather          <tt>true</tt> if the send is for a gather,
     *                          otherwise <tt>false</tt>
     * 
     * @throws  IOException     if there is any communication error while
     *                          sending
     */
    private synchronized static void send( int rank, Object o, int tag,
            boolean gather ) throws IOException {
    Node node = nodes[ rank ];
        if ( gather ) {
            node.out.writeByte( Node.Command.GATHER );
        } else {
            node.out.writeByte( Node.Command.RECEIVE );
        }
        node.out.writeInt( tag );
        node.out.flush();
        
        ObjectOutputStream oos = new ObjectOutputStream( node.out );
        oos.writeObject( o );
    }
    
    /**
     * Receives an object from the specified node under the tag of <tt>0</tt>.
     * This method will block until an object is available.
     * 
     * @param   rank                    the rank of the node to receive the
     *                                  object from
     * 
     * @return                          the object from the specified node with
     *                                  the tag of <tt>0</tt>
     * 
     * @throws  InterruptedException    if the thread is interrupted while
     *                                  waiting for an object
     */
    public static Object receive( int rank ) throws InterruptedException {
        return receive( rank, 0 );
    }
    
    /**
     * Receives an object from the specified node under the specified tag.  This
     * method will block until an object is available.
     * 
     * @param   rank                    the rank of the node to receive the
     *                                  object from
     * @param   tag                     the tag under which the object is
     *                                  associated
     * 
     * @return                          the object from the specified node with
     *                                  the specified tag
     * 
     * @throws  InterruptedException    if the thread is interrupted while
     *                                  waiting for an object
     */
    public static Object receive( int rank, int tag )
            throws InterruptedException {
        return receive( rank, tag, false );
    }
    
    /**
     * Helper method for receiving an object for either a receive or a gather.
     * 
     * @param   rank                    the rank of the node to receive the
     *                                  object from
     * @param   tag                     the tag under which the object is
     *                                  associated
     * @param   gather                  <tt>true</tt> if the object is being
     *                                  received as a gather, otherwise
     *                                  <tt>false</tt>
     * 
     * @return                          the object from the specified node with
     *                                  the specified tag
     * 
     * @throws  InterruptedException    if the thread is interrupted while
     *                                  waiting for an object
     */
    public static Object receive( int rank, int tag, boolean gather )
            throws InterruptedException {
        Object obj = null;
        if ( gather ) {
            obj = ( nodes[ rank ].gathered ).take( tag );
        } else {
            obj = ( nodes[ rank ].received ).take( tag );
        }
        return obj;
    }
    
    /**
     * Performs an MPI <b>gather</b> operation with the specified root, and
     * submits the specified object for this particular node. The gather is
     * performed under the tag of <tt>0</tt>.
     * 
     * @param   root                    the root node of the gather
     * @param   o                       the object to be submitted for this node
     * 
     * @return                          an array of size <tt>Comm.size()</tt>
     *                                  with the objects from all nodes on the
     *                                  root node; non-root nodes return
     *                                  <tt>null</tt>
     * 
     * @throws  IOException             if there is any communication error
     * @throws  InterruptedException    if there is any error while waiting for
     *                                  this method to finish
     */
    public static Object[] gather( int root, Object o ) throws IOException,
            InterruptedException {
        return gather( root, o, 0 );
    }
    /**
     * Performs an MPI <b>gather</b> operation with the specified root, and
     * submits the specified object for this particular node.  The gather is
     * performed under the specified tag.
     * 
     * @param   root                    the root node of the gather
     * @param   o                       the object to be submitted for this node
     * @param   tag                     the tag under which to associate the
     *                                  gather
     * 
     * @return                          an array of size <tt>Comm.size()</tt>
     *                                  with the objects from all nodes on the
     *                                  root node; non-root nodes return
     *                                  <tt>null</tt>
     * 
     * @throws  IOException             if there is any communication error
     * @throws  InterruptedException    if there is any error while waiting for
     *                                  this method to finish
     */
    public static Object[] gather( int root, Object o, int tag )
            throws IOException, InterruptedException {
        if ( rank != root ) {
            send( root, o, tag, true );
            return null;
        } else {
            Object[] arr = new Object[ size ];
            // wait for all objects before returning
            for ( int i = 0; i < size; i++ ) {
                if ( i != rank ) {
                    arr[ i ] = receive( i, tag, true );
                } else {
                    arr[ i ] = o;
                }
            }
            return arr;
        }
    }
    
    /**
     * Performs an MPI <b>all-gather</b> operation, and submits the specified
     * object for this particular node. The all-gather is performed under the
     * tag of <tt>0</tt>.
     * 
     * @param   o                       the object to be submitted for this node
     * 
     * @return                          an array of size <tt>Comm.size()</tt>
     *                                  with the objects from all nodes
     * 
     * @throws  IOException             if there is any communication error
     * @throws  InterruptedException    if there is any error while waiting for
     *                                  this method to finish
     */
    public static Object[] allGather( Object o ) throws IOException,
            InterruptedException {
        return allGather( o, 0 );
    }
    
    /**
     * Performs an MPI <b>all-gather</b> operation, and submits the specified
     * object for this particular node. The all-gather is performed under the
     * specified tag.
     * 
     * @param   o                       the object to be submitted for this node
     * @param   tag                     the tag under which to associate the
     *                                  all-gather
     * 
     * @return                          an array of size <tt>Comm.size()</tt>
     *                                  with the objects from all nodes
     * 
     * @throws  IOException             if there is any communication error
     * @throws  InterruptedException    if there is any error while waiting for
     *                                  this method to finish
     */
    public static Object[] allGather( Object o, int tag ) throws IOException,
            InterruptedException {
        // send all nodes my object
        for ( int i = 0; i < size; i++ ) {
            if ( i != rank ) {
                send( i, o, tag, true );
            }
        }
        
        // wait for all objects before returning
        Object[] arr = new Object[ size ];
        for ( int i = 0; i < size; i++ ) {
            if ( i != rank ) {
                arr[ i ] = receive( i, tag, true );
            } else {
                arr[ i ] = o;
            }
        }
        return arr;
    }
    
    /**
     * Shuts down all node connections from this communicator.
     */
    protected static void shutdown() {
        // close all node connections
        for ( Node node : nodes ) {
            if ( node != null ) {
                try {
                    node.socket.close();
                } catch ( Exception e ) {
                    System.err.println( "Error closing socket." );
                    e.printStackTrace( System.err );
                }
            }
        }
    }
    
    /**
     * This private inner thread represents a remote node, and handles all input
     * and output with the node it represents.
     * 
     * @author  Peter O. Erickson
     * 
     * @version $Id: Comm.java 27 2009-04-17 02:30:39Z PErickson87 $
     */
    protected static final class Node extends Thread {
        
        /** The rank of this node. */
        private short rank;
        
        /** The socket connected to the node. */
        private Socket socket;
        
        /** The input stream from this node. */
        private DataInputStream in;
        
        /** The output stream for this node. */
        private DataOutputStream out;
        
        /** A queue of received objects from this node. */
        private MessageQueue received;
        
        /** A queue of gathered objects from this node. */
        private MessageQueue gathered;
        
        /**
         * Holds the valid set of commands for the control layer.
         * 
         * @author  Peter O. Erickson
         * 
         * @version $Id: Comm.java 27 2009-04-17 02:30:39Z PErickson87 $
         */
        protected static final class Command {
            
            /** Byte signal to indicate a barrier. */
            private static final byte BARRIER = 0x2;
            
            /** Byte signal to indicate that an object is to be received. */
            private static final byte RECEIVE = 0x1;
            
            /** Byte signal to indicate that an object is to be gathered. */
            private static final byte GATHER = 0x3;
            
            /** Byte signal to indicate the end of the node's life. */
            protected static final byte NODE_END = 0x9;
            
        } // Commmand
        
        /**
         * Constructs a new node proxy.
         * 
         * @param   rank            the rank of the node
         * @param   socket          the socket connected to this node
         * 
         * @throws  IOException     if there is any fatal IO error
         */
        protected Node( short rank, Socket socket ) throws IOException {
            this.rank = rank;
            this.socket = socket;
            
            this.out = new DataOutputStream( socket.getOutputStream() );
            this.in = new DataInputStream( socket.getInputStream() );
            
            received = new MessageQueue();
            gathered = new MessageQueue();
            
            // run the listener thread
            this.start();
        }
        
        /**
         * Runs the listener for data from the actual node.
         */
        @Override
        public void run() {
            byte cmd = 0;
            while ( !( socket.isClosed() ) ) {
                try {
                    // read the incoming command
                    cmd = in.readByte();
                    
                    switch ( cmd ) {
                        // contribute to an existing barrier, or start new
                        case Command.BARRIER: {
                            try {
                                 barrier.await();
                            } catch ( Exception e ) {
                                System.err.println(
                                        "Error waiting node at barrier." );
                                e.printStackTrace( System.err );
                            }
                        } break;
                        
                        // receive an object from this node, and cache it
                        case Command.RECEIVE: {
                            try {
                                int tag = in.readInt();
                                
                                // use a custom ObjectInputStream with the
                                // Comm's network classloader
                                ObjectInputStream ois =
                                        new ClassLoaderOIS( in, ncl );
                                Object o = ois.readObject();
                                
                                received.add( o, tag );
                            } catch ( Exception e ) {
                                // this shouldn't happen with a network class-
                                // loader
                                System.err.println(
                                        "Network classloader error." );
                                e.printStackTrace( System.err );
                            }
                        } break;
                        
                        // receive an object from this node, and cache it
                        case Command.GATHER: {
                            try {
                                int tag = in.readInt();
                                
                                // use a custom ObjectInputStream with the
                                // Comm's network classloader
                                ObjectInputStream ois =
                                        new ClassLoaderOIS( in, ncl );
                                Object o = ois.readObject();
                                
                                gathered.add( o, tag );
                            } catch ( Exception e ) {
                                // this shouldn't happen with a network class-
                                // loader
                                System.err.println(
                                        "Network classloader error." );
                                e.printStackTrace( System.err );
                            }
                        } break;
                        
                        // end the node and call the end barrier
                        case Command.NODE_END: {
                            break;
                        }
                    }
                    
                } catch ( EOFException e ) {
                    // kill the node loop
                    break;
                } catch ( IOException e ) {
                    // likely something broke, kill the loop
                    break;
                }
            }
            
            // end the network class loader server if we're the root node
            if ( Comm.rank == 0 ) {
                // wait at the end barrier
                try {
                    endBarrier.await();
                } catch ( Exception e ) {
                    System.err.println(
                            "Error waiting at end barrier." );
                e.printStackTrace( System.err );
                }
            }
        }
        
        /**
         * A blocking message queue that stores received objects by tag.
         * 
         * @author  Peter O. Erickson
         * 
         * @version $Id: Comm.java 27 2009-04-17 02:30:39Z PErickson87 $
         */
        private static final class MessageQueue
                extends HashMap<Integer, LinkedBlockingQueue> {
            
            /**
             * Adds an object to the message queue under the specified tag.
             * 
             * @param   o       the object to be added to the queue
             * @param   tag     the tag under which to associate the object
             */
            void add( Object o, int tag ) {
                q( tag ).offer( o );
            }
            
            /**
             * Take an object with the specified tag from the queue, or block
             * until one becomes available.
             * 
             * @param   tag                     the tag of the object to be
             *                                  taken out of the queue
             * 
             * @return                          an object with the specified tag
             * 
             * @throws  InterruptedException    if this thread is interrupted
             *                                  while waiting for an object to
             *                                  become available
             */
            Object take( int tag ) throws InterruptedException {
                return q( tag ).take();
            }
            
            /**
             * Returns the queue associated with the specified tag, or creates
             * one if no such queue exists.
             * 
             * @param   tag     the tag of the queue to be returned
             * 
             * @return          the queue with the specified tag
             */
            private synchronized BlockingQueue q( int tag ) {
                LinkedBlockingQueue q;
                if ( ( q = get( tag ) ) == null ) {
                    put( tag, q = new LinkedBlockingQueue() );
                }
                return q;
            }
            
        } // MessageQueue
        
    } // Node
    
} // Comm
